#!/usr/bin/env/ python
# -*- coding:utf-8 -*-

import os
import sys
import traceback
import configparser
import json
import platform
import subprocess

from utils.log import logger
from mydb import MyDB
from clickhouse_cli import ClickhouseCli
from postgre_sql_cli import PostgreSQLCli
class DatabaseChecker(object):
    def __init__(self):
        try:
            config = configparser.ConfigParser()
            hosts_config_file_path = 'conf/hosts.conf'
            databases_config_file_path = 'conf/databases.conf'
            db_type_config_file_path = 'conf/database_types.conf'
            self.table_total_row_count_check_conf_file = 'conf/table_record_check.conf'
            if os.getenv('DB_CHECKER_DIR'):
                hosts_config_file_path = os.path.join(os.environ['DB_CHECKER_DIR'], hosts_config_file_path)
                db_type_config_file_path = os.path.join(os.environ['DB_CHECKER_DIR'], db_type_config_file_path)
                databases_config_file_path = os.path.join(os.environ['DB_CHECKER_DIR'], databases_config_file_path)
                self.table_total_row_count_check_conf_file = os.path.join(os.environ['DB_CHECKER_DIR'], self.table_total_row_count_check_conf_file)
            
            with open(self.table_total_row_count_check_conf_file, 'r', encoding='utf-8') as f:
                self.table_record_check_config = json.loads(f.read())
                        
            config.read(databases_config_file_path, encoding='utf-8')
            self.databases_config = self.parse_databases_conf(config)
        except Exception:
            logger.error('程序初始化操作失败：%s' % traceback.format_exc())
            sys.exit()
    
    
    def check_host_connectivity(self, host):
        '''主机连通性检查'''

        args = f'ping {host} -n 2'
        if platform.system() == 'Linux':
            args =  f'ping {host} -c 2'

        with subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True,
                              universal_newlines=True) as proc:
            try:
                outs, errs = proc.communicate(timeout=10)
                if 'Ping 请求找不到主机' in outs or 'Name or service not known' in outs:
                    logger.error(f'执行命令{args}失败：输出：\n{outs}')
                    return False
                elif '已发送 = 2，已接收 = 2，丢失 = 0 (0% 丢失)' in outs or '2 packets transmitted, 2 received, 0% packet loss' in outs:
                    logger.debug(f'执行命令{args}成功：输出：\n{outs}')
                    return True
                else:
                    logger.error(f'执行命令{args}失败：输出：\n{outs}')
                    return False
            except subprocess.TimeoutExpired:
                proc.kill()
                logger.error(f'执行命令{args}失败:\n{traceback.format_exc()}')
                return False
    
    def parse_databases_conf(self, dabatases_config):
        '''解析自定义数据库配置'''
        
        result = {}  # 存放需要访问的目标机器配置
        temp_dict = {}
        dbs_excluded = {}
        for section in dabatases_config.sections():
            db_config = {}
            option_list = dabatases_config.options(section)
            
            for option in option_list:
                db_config[option] = dabatases_config[section][option]
                
            db_host = db_config.get('db_host')
            if not self.check_host_connectivity(db_host):
                logger.warn(f'数据库主机ping不通，忽略节点({section})配置')
                continue
            db_port = db_config.get('db_port')
            key = f'{db_host}-{db_port}'
            if key not in temp_dict:
                temp_dict[key] = 1
            else:
                logger.warn(f'数据库 主机-端口 组合值不唯一，配置节点：{section}')
                continue
            
            databases_excluded = db_config.get('databases_excluded')
            if databases_excluded:
                databases_excluded = databases_excluded.replace('，', ',').replace(' ', '').rstrip(',').split(',')
            else:
                databases_excluded = []
            
            schemas_excluded = db_config.get('schemas_excluded')
            if schemas_excluded:
                schemas_excluded = schemas_excluded.replace('，', ',').replace(' ', '').rstrip(',').split(',')
            else:
                schemas_excluded = []
                
            db_type = db_config.get('db_type')
            if db_type not in result:
                result[db_type] = {'db_info':[], 'databases_excluded':[], 'schemas_excluded':[]}
            result[db_type]['db_info'].append(db_config)
            
            for item in databases_excluded:
                if item not in result[db_type]['databases_excluded']:
                    result[db_type]['databases_excluded'].append(item)
            
            for item in schemas_excluded:
                if item not in result[db_type]['schemas_excluded']:
                    result[db_type]['schemas_excluded'].append(item)
                    
        return result
    

    def analyze_data_collected(self, data_collected, db_type):
        '''分析采集的数据(同类数据库比较)
        @params data_collected:  非opengauss {'host:port': {'database_name': {'table_name': {'total_record_num':total_record_num, sample_records: [{record1},{record2}]}}}}
        opengauss {'host:port': {'database_name': {'schema_name': {'table_name': {'total_record_num':total_record_num, sample_records: [{record1},{record2}]}}}}}
        '''
 
        result = {'databases_missing': [], 'tables_missing':[], 'tables_record_num_abnormal': {}}
        
        all_dbs = set()
        all_schemas = {}
        all_tables = {}
        table_total_record_num_dict = {}
        table_total_record_num_info = {}
        
        if db_type != 'opengauss':
            for host_port, data in data_collected.items():
                for db, table_info in data.items():
                    all_dbs.add(db)
                    
                    if db not in all_tables:
                        all_tables[db] = set()
                    if db not in table_total_record_num_dict:
                        table_total_record_num_dict[db] = {}
                        table_total_record_num_info[db] = {}
                    for table, _data in table_info.items():
                        all_tables[db].add(table)
                        if table not in table_total_record_num_dict[db]:
                            table_total_record_num_dict[db][table] = []
                            table_total_record_num_info[db][table] = []
                        total_record_num = _data.get('total_record_num')
                        table_total_record_num_dict[db][table].append(total_record_num)
                        table_total_record_num_info[db][table].append(f'实例：{host_port} 数据库：{db} 表：{table}, 记录总数：{total_record_num}')
            
            # 比较库
            for host_port, data in data_collected.items():
                host_dbs = set(data.keys())
                missing_dbs = all_dbs - host_dbs
                if missing_dbs:
                    result['databases_missing'].append(f'实例{host_port} 缺失 {db_type} 数据库: {missing_dbs} ')
                    
            # 比较表
            for host_port, data in data_collected.items():
                for db, table_info in data.items():
                    db_tales = set(table_info.keys())
                    missing_tables = all_tables.get(db) - db_tales
                    if missing_tables:
                        result['tables_missing'].append(
                            f'实例{host_port} {db_type} 数据库 {db} 缺失表：{missing_tables}')
            
            # 比较表记录数
            for db, table_info in table_total_record_num_dict.items():
                for table, total_record_nums in table_info.items():
                    if len(set(total_record_nums)) > 1:
                        if db not in result['tables_record_num_abnormal']:
                            result['tables_record_num_abnormal'][db] = {}
                        result['tables_record_num_abnormal'][db][table] = table_total_record_num_info.get(db, {}).get(
                            table)
            return result
        else:
            for host_port, data in data_collected.items():
                for db, schema_info in data.items():
                    all_dbs.add(db)
                    
                    if db not in all_tables:
                        all_tables[db] = {}
                        all_schemas[db] = set()
                    if db not in table_total_record_num_dict:
                        table_total_record_num_dict[db] = {}
                        table_total_record_num_info[db] = {}
                    
                    for schema, table_info in schema_info.items():
                        all_schemas[db].add(schema)
                        if schema not in table_total_record_num_dict[db]:
                            table_total_record_num_dict[db][schema] = {}
                            table_total_record_num_info[db][schema] = {}
                            all_tables[db][schema] = set()
                            
                        for table, _data in table_info.items():
                            all_tables[db][schema].add(table)
                            if table not in table_total_record_num_dict[db][schema]:
                                table_total_record_num_dict[db][schema][table] = []
                                table_total_record_num_info[db][schema][table] = []
                            total_record_num = _data.get('total_record_num')
                            table_total_record_num_dict[db][schema][table].append(total_record_num)
                            table_total_record_num_info[db][schema][table].append(
                                f'实例：{host_port} 数据库：{db} 模式：{schema} 表：{table}, 记录总数：{total_record_num}')
                        
            # 比较库
            for host_port, data in data_collected.items():
                host_dbs = set(data.keys())
                missing_dbs = all_dbs - host_dbs
                if missing_dbs:
                    result['databases_missing'].append(f'实例{host_port} 缺失 {db_type} 数据库: {missing_dbs} ')
                    
            # 比较模式
            for host_port, data in data_collected.items():
                for db, schema_info in data.items():
                    db_schemas = set(schema_info.keys())
                    missing_schemas = all_schemas.get(db) - db_schemas
                    if missing_schemas:
                        result['schema_missing'].append(f'实例{host_port} {db_type} 数据库 {db} 缺失模式：{missing_schemas}')
                        
            # 比较表
            for host_port, data in data_collected.items():
                for db, schema_info in data.items():
                    for schema, table_info in schema_info.items():
                        db_tales = set(table_info.keys())
                        missing_tables = all_tables.get(db).get(schema) - db_tales
                        if missing_tables:
                            result['tables_missing'].append(f'实例{host_port} {db_type} 数据库 {db} 模式 {schema} 缺失表：{missing_tables}')
                            
            # 比较表记录数
            for db, schema_info in table_total_record_num_dict.items():
                for schema, table_info in schema_info.items():
                     for table, total_record_nums in table_info.items():
                        if len(set(total_record_nums)) > 1:
                            key = f'{db}_{schema}'
                            if key not in result['tables_record_num_abnormal']:
                                result['tables_record_num_abnormal'][key] = {}
                            result['tables_record_num_abnormal'][key][table] = table_total_record_num_info.get(db, {}).get(schema, {}).get(table)
            return result

    
    def check_mariadb(self, db_configs, db_type):
        '''检查mariadb数据库'''
        
        data_collected = {}  # {'host:port': {'database_name': {'table_name': {'total_record_num':total_record_num, sample_records: [{record1},{record2}]}}}}
        logger.info(f'正在检查{db_type}数据库是否同步')
        for db_config in db_configs.get('db_info', []):
            db_host = db_config.get('db_host')
            db_port = db_config.get('db_port')
            
            db_obj = None
            try:
                db_obj = MyDB(db_name = db_config.get('db_name'),
                              db_host = db_host,
                              port = int(db_port),
                              user = db_config.get('db_user'),
                              password = db_config.get('db_passwd'),
                              charset = db_config.get('db_charset')
                              )
                
                host_port = f'{db_host}:{db_port}'
                if host_port not in data_collected:
                    data_collected[host_port] = {}
                
                result = db_obj.select_many('SHOW DATABASES')
                databases_excluded = db_configs.get('databases_excluded')
                for item in result:
                    db_name = item[0]
                    if db_name in databases_excluded:
                        continue
                    if db_name not in data_collected[host_port]:
                        data_collected[host_port][db_name] = {}
                    
                    db_obj.select(f'USE `{db_name}`')
                    tables = db_obj.select(f'SHOW TABLES')
                    for item in tables:
                        table = item[0]
                        if table not in data_collected[host_port][db_name]:
                            data_collected[host_port][db_name][table] = {}
                        query = f'SELECT COUNT(1) AS total_record_num FROM `{db_name}`.`{table}`'
                        key = f'{db_type}.{db_name}.{table}'
                        query = self.table_record_check_config.get(key, {}).get('total_record_num', query)
                        res = db_obj.select(query, dictionary=True)
                        data_collected[host_port][db_name][table]['total_record_num'] = res.get('total_record_num', 0)
                        
                        query = self.table_record_check_config.get(key, {}).get('sample_records')
                        if query:
                            res = db_obj.select_many(query, dictionary=True)
                            data_collected[host_port][db_name][table]['sample_records'] = res
            except Exception:
                logger.error('%s' % traceback.format_exc())
            finally:
                if db_obj:
                    db_obj.close()
        res = self.analyze_data_collected(data_collected, db_type)
        self.print_abnormal_messages(res, db_type)
    
    def check_clickhouse(self, db_configs, db_type):
        '''检查clickhouse数据库'''
        
        data_collected = {}  # {'host:port': {'database_name': {'table_name': {'total_record_num':total_record_num, sample_records: [{record1},{record2}]}}}}
        logger.info(f'正在检查{db_type}数据库是否同步')
        for db_config in db_configs.get('db_info', []):
            db_host = db_config.get('db_host')
            db_port = db_config.get('db_port')
            
            db_obj = None
            try:
                db_obj = ClickhouseCli(db_name=db_config.get('db_name'),
                              db_host=db_host,
                              port=int(db_port),
                              user=db_config.get('db_user'),
                              password=db_config.get('db_passwd'),
                              connect_timeout=int(db_config.get('connect_timeout') or '15')
                              )
                
                host_port = f'{db_host}:{db_port}'
                if host_port not in data_collected:
                    data_collected[host_port] = {}
                
                result = db_obj.select_many('SHOW DATABASES')
                databases_excluded = db_configs.get('databases_excluded')
                for item in result:
                    db_name = item[0]
                    if db_name in databases_excluded:
                        continue
                    if db_name not in data_collected[host_port]:
                        data_collected[host_port][db_name] = {}

                    tables = db_obj.select(f'SHOW TABLES FROM `{db_name}`')
                    for table in tables:
                        if table not in data_collected[host_port][db_name]:
                            data_collected[host_port][db_name][table] = {}
                        query = f'SELECT COUNT(1) AS total_record_num FROM `{db_name}`.`{table}`'
                        key = f'{db_type}.{db_name}.{table}'
                        query = self.table_record_check_config.get(key, {}).get('total_record_num', query)
                        res = db_obj.select(query)
                        data_collected[host_port][db_name][table]['total_record_num'] = res[0]
                        
                        query = self.table_record_check_config.get(key, {}).get('sample_records')
                        if query:
                            res = db_obj.select_many(query)
                            data_collected[host_port][db_name][table]['sample_records'] = res
            except Exception:
                logger.error('%s' % traceback.format_exc())
            finally:
                if db_obj:
                    db_obj.close()
        
        res = self.analyze_data_collected(data_collected, db_type)
        self.print_abnormal_messages(res, db_type)
    
    def check_opengauss(self, db_configs, db_type):
        '''检查opengauss数据库'''
        
        data_collected = {}  # {'host:port': {'database_name': {'table_name': {'total_record_num':total_record_num, sample_records: [{record1},{record2}]}}}}
        logger.info(f'正在检查{db_type}数据库是否同步')
        for db_config in db_configs.get('db_info', []):
            db_host = db_config.get('db_host')
            db_port = db_config.get('db_port')
            
            db_obj = None
            try:
                db_obj = PostgreSQLCli(db_name=db_config.get('db_name'),
                                       db_host=db_host,
                                       port=int(db_port),
                                       user=db_config.get('db_user'),
                                       password=db_config.get('db_passwd')
                                       )

                host_port = f'{db_host}:{db_port}'
                if host_port not in data_collected:
                    data_collected[host_port] = {}
                
                result = db_obj.select_many('SELECT datname AS "Database" FROM pg_database;')
                databases_excluded = db_configs.get('databases_excluded')
                table_schemas_excluded = db_configs.get('schemas_excluded')
                table_schemas_excluded_temp = ["'%s'" % item for item in table_schemas_excluded]
                table_schemas_excluded_str = ','.join(table_schemas_excluded_temp)
                for item in result:
                    db_name = item[0]
                    if db_name in databases_excluded:
                        continue
                    if db_name not in data_collected[host_port]:
                        data_collected[host_port][db_name] = {}
                    
                    db_obj = PostgreSQLCli(db_name=db_config.get('db_name'),
                                           db_host=db_host,
                                           port=int(db_port),
                                           user=db_config.get('db_user'),
                                           password=db_config.get('db_passwd')
                                           )
                   
                    query_table_with_schema_sql = f'''SELECT table_schema AS schema_name, table_name
                                                     FROM information_schema.tables
                                                     WHERE table_type = 'BASE TABLE'
                                                     AND table_schema NOT IN ({table_schemas_excluded_str})'''
                    res = db_obj.select_many(query_table_with_schema_sql)
                    
                    for schema_name, table in res:
                        if f'{db_name}.{schema_name}' in table_schemas_excluded:
                            continue
                            
                        if schema_name not in data_collected[host_port][db_name]:
                            data_collected[host_port][db_name][schema_name] = {}
                            
                        if table not in data_collected[host_port][db_name][schema_name]:
                            data_collected[host_port][db_name][schema_name][table] = {}
                        query = f'SELECT COUNT(1) AS total_record_num FROM {db_name}.{schema_name}."{table}"' # 表名加双引号，防止表名中存在特殊字符比如 / 时报错
                        key = f'{db_type}.{db_name}.{schema_name}.{table}'
                        query = self.table_record_check_config.get(key, {}).get('total_record_num', query)
                        res = db_obj.select(query)
                        data_collected[host_port][db_name][schema_name][table]['total_record_num'] = res[0]

                        query = self.table_record_check_config.get(key, {}).get('sample_records')
                        if query:
                            res = db_obj.select_many(query)
                            data_collected[host_port][db_name][schema_name][table]['sample_records'] = res
            except Exception:
                logger.error('%s' % traceback.format_exc())
            finally:
                if db_obj:
                    db_obj.close()
        
        res = self.analyze_data_collected(data_collected, db_type)
        self.print_abnormal_messages(res, db_type)
        
    def print_abnormal_messages(self, result_data, db_type):
        '''打印异常信息'''
        
        status_normal = True
        databases_missing = result_data.get('databases_missing')
        if databases_missing:
            logger.info('缺库异常：\n' + '\n'.join(databases_missing))
            status_normal = False
        
        schema_missing = result_data.get('schema_missing')
        if schema_missing:
            logger.info('缺模式异常：\n' + '\n'.join(databases_missing))
            status_normal = False
        
        tables_missing = result_data.get('tables_missing')
        if tables_missing:
            logger.info('缺表异常：\n' +'\n'.join(tables_missing))
            status_normal = False
        
        for key, table_abnormal_info in result_data.get('tables_record_num_abnormal').items():
            for table, record_num_abnormal_info in table_abnormal_info.items():
                logger.info(f'{table}表记录总数不一致异常：\n' + '\n'.join(record_num_abnormal_info))
                status_normal = False
        
        if status_normal:
            logger.info(f'{db_type}数据库同步检查完毕，未发现异常，请检查查程序运行是否报错')
    
    def check_databases(self):
        '''检查数据库'''

        db_type_check_func_map = {'mariadb': self.check_mariadb,
                                  'mysql': self.check_mariadb,
                                  'clickhouse': self.check_clickhouse,
                                  'opengauss': self.check_opengauss,
                                  'postgres': self.check_opengauss
                                  }
                            
        if not self.databases_config:
            logger.warn('获取数据库配置为空，请检查配置：conf/hosts.conf、database_checker.conf、database_types.conf')
            sys.exit(1)
        
        for db_type, db_configs in self.databases_config.items():
            if len(db_configs) < 2: # 同类型数据库对比，少于2个实例
                logger.warn(f'{db_type}数据库类型只包含一个实例，无需比较，已跳过')
                continue
            check_func = db_type_check_func_map.get(db_type.lower())
            if check_func:
                check_func(db_configs, db_type)
               
    def run(self):
        self.check_databases()
 


